import tensorflow as tf 
import sklearn
from sklearn import datasets, model_selection
import os 
import numpy as np 

# 加载和划分数据集
# sklearn加载的是numpy格式
housing = datasets.fetch_california_housing()
train_x,test_x,train_y,test_y = model_selection.train_test_split(
    housing.data,housing.target,random_state=21)
train_x,valid_x,train_y,valid_y = model_selection.train_test_split(
    train_x,train_y,random_state=42)

# 均值方差归一化
scaler = sklearn.preprocessing.StandardScaler()
train_x = scaler.fit_transform(train_x)
valid_x = scaler.transform(valid_x)
test_x = scaler.transform(test_x)

train_data = np.c_[train_x,train_y]  # 合并
valid_data = np.c_[valid_x,valid_y]
test_data = np.c_[test_x,test_y]
header = housing.feature_names+['MidianHouseValue']
header = ",".join(header)
# print(header)  # MedInc,HouseAge,AveRooms,AveBedrms,Population,AveOccup,Latitude,Longitude,MidianHouseValue

output_dir = os.path.join(os.getcwd(),'tf_data','generate')
if not os.path.exists(output_dir):
    os.mkdir(output_dir)

def save_to_csv(output_dir,data,name_prefix,header=None,n_parts=10):
    path_format = os.path.join(output_dir,"{}_{:02d}.csv")  # 两个数字的整数  不足两个前面补0
    filenames = [ ]

    for file_idx, row_idx  in enumerate(np.array_split(np.arange(len(data)),n_parts)): # 把data的每个数据编号 并分成n_parts份
        part_csv = path_format.format(name_prefix,file_idx)  # 给每一份文件命名
        filenames.append(part_csv)
        with open(part_csv,'wt',encoding='utf-8') as f:  # 打开文件写入内容
            if header is not None:
                f.write(header + '\n')
            for row in row_idx:
                f.write(",".join([repr(element) for element in data[row]]) + '\n')  # 取出一行的数据 用逗号分隔(csv格式) 然后写入 
    return filenames

train_filenames = save_to_csv(output_dir,train_data,"train",header,n_parts=20)
valid_filenames = save_to_csv(output_dir,valid_data,"valid",header,n_parts=10)
test_filenames = save_to_csv(output_dir,test_data,"test",header,n_parts=10)

print(train_filenames)
print(valid_filenames)
print(test_filenames)
